import argparse
import os

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Data converter arg parser')
    parser.add_argument(
        '--load_dir',
        type=str,
        default='/workspace/data/blender/mesh_data')
    parser.add_argument(
        '--mesh_type',
        type=str,
        default=None)
    parser.add_argument(
        '--save_root',
        type=str,
        default='/workspace/data/blender/render_cmd_sh/')
    parser.add_argument(
        '--batch_size',
        type=int,
        default=5)
    args = parser.parse_args()

    for cmd_type in ['debug', 'render', 'batch']:
        
        if cmd_type == 'render' or cmd_type == 'batch':
            postfix = ' --debug_mode=-1'
        else:
            postfix = '--save_root_dir=render_res/debug'

        for texture_name in ['mesh', 'texture']:

            with open(os.path.join(args.save_root, cmd_type, texture_name + '_cmds.sh'), 'w') as f:

                for mesh_type in os.listdir(args.load_dir):

                    if args.mesh_type is not None and mesh_type != args.mesh_type:
                        continue
                    f.write(f"\n# mesh_type: {mesh_type}\n")
                    mesh_list = os.listdir(os.path.join(args.load_dir, mesh_type))
                    texture_postfix = '' if texture_name == 'mesh' else '--is_texture '
                    
                    if cmd_type == 'batch':
                        for i in range(8):
                            f.write(f"python mesh.py --mesh_type {mesh_type} --mesh_list " + ' '.join(mesh_list[i * args.batch_size:(i + 1) * args.batch_size]) + f" --cuda_id {i} " + texture_postfix + postfix + "\n")
                            # ::8]) + f" --cuda_id {i} " + texture_postfix + postfix + "\n")
    
                    else:
                        for idx, mesh in enumerate(mesh_list):
                            f.write(f"python mesh.py --mesh_type {mesh_type} --mesh_list {mesh} --cuda_id {idx % 8} " + texture_postfix + postfix + "\n")